
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.nio.file.Paths;
import java.security.SecureRandom;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AesCryption {
	public static boolean debug_mode = false;
	public static String file_to_save = "";
	/**
	 * 输入2个参数：待加密文件和输出文件夹，均使用绝对路径。
	 * 
	 * @param args
	 * @throws Exception
	 */
	public static void main(String[] args) throws Exception {
		File in_file = new File(args.length > 0 ? args[0] : "key.png");
		if (!in_file.exists()) {
			System.out.println("Quit. Error. The input file \"" + in_file.getAbsolutePath() + "\" does not exist.");
			return;
		}

		// output files:
		// 1. encrypt: [out_dir]/aes_files/filename.ext.tmp
		// 2. decrypt: [out_dir]/aes_files/filename.ext
		// 3. result: [out_dir]/result.csv
		File out_dir = new File(args.length > 1 ? args[1] : "output");
		if (!out_dir.exists() && !out_dir.mkdir()) {
			System.out.println("Quit. Error. The directory '" + out_dir.getAbsolutePath() + "' is failed to create.");
			return;
		}
		
		Common.file_to_save = Paths.get(out_dir.getAbsolutePath(), "result.csv").toString();
		File files_dir = new File(Paths.get(out_dir.getAbsolutePath(), "aes_files").toString());
		files_dir.mkdir();
		File temp_file = new File(Paths.get(files_dir.getAbsolutePath(), in_file.getName() + ".tmp").toString());
		File output_file = new File(Paths.get(files_dir.getAbsolutePath(), in_file.getName()).toString());

		if (debug_mode) 
		{
			System.out.print("------------------ Arguments Info -----------------\n");
			System.out.println("input file:    " + in_file.getAbsolutePath());
			System.out.println("encrypted:     " + temp_file.getAbsolutePath());
			System.out.println("decrypted:     " + output_file.getAbsolutePath());
			System.out.print("---------------------------------------------------\n");
			return;
		}

		AesCryption ac = new AesCryption();
		for (int i = 0; i < 11; i++) {
			ac.process(in_file.getAbsolutePath(), temp_file.getAbsolutePath(), true, i);
		}
		Thread.sleep(100);
		for (int i = 0; i < 11; i++) {
			ac.process(temp_file.getAbsolutePath(), output_file.getAbsolutePath(), false, i);
		}
	}

	public static void fileTest() throws Exception {
		String input_file = "D:\\data\\title_text.txt";
		input_file = "D:\\data\\bts.bin";
		String encrypt_file = input_file + ".tmp";
		String output_file = input_file + input_file.substring(input_file.length() - 4);

		AesCryption ac = new AesCryption();
		ac.process(input_file, encrypt_file, true, 0);
		ac.process(encrypt_file, output_file, false, 0);

		System.out.println("Done.");
	}

	byte[] IV;
	SecretKey secKey;

	public AesCryption() {
		try {
			KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
			keyGenerator.init(128);
			secKey = keyGenerator.generateKey();
			IV = new byte[16];
			SecureRandom random = new SecureRandom();
			random.nextBytes(IV);
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public void process(String infile, String outfile, boolean encrypt, int test_id) throws Exception {
		// 初始化所有相关变量
		FileInputStream fin = new FileInputStream(infile);
		FileOutputStream fout = new FileOutputStream(outfile);
		int filesize = fin.available();
		byte[] bts = new byte[encrypt ? 1024576 : 1024592];
		double t1 = 0, t2 = 0, t3 = 0;
		long timeTick = 0;
		int readCount = bts.length;
		Cipher cipher = getCipher(encrypt ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE);

		// 加解密过程包括读取、加解密和写回文件。
		while (true) {
			// t1: 文件读取时间
			timeTick = System.nanoTime();
			readCount = fin.read(bts, 0, bts.length);
			if (readCount <= 0)
				break;
			t1 += System.nanoTime() - timeTick;

			// t2: 文件加密时间
			timeTick = System.nanoTime();
			if (readCount < bts.length) {
				byte[] bts1 = new byte[readCount];
				for (int i = 0; i < bts1.length; i++)
					bts1[i] = bts[i];
				bts = bts1;
			}
			byte[] result = cipher.doFinal(bts);
			t2 += System.nanoTime() - timeTick;

			// t3: 文件写入时间
			timeTick = System.nanoTime();
			fout.write(result, 0, result.length);
			t3 += System.nanoTime() - timeTick;
		}
		fin.close();
		fout.close();

		String type = encrypt ? "aes_en" : "aes_de";
		// 4个过程使用统一的时间，方便归类
		String datetime = Common.getTimeString();
		Common.appendToFile(datetime, test_id, type, "load", filesize, t1, infile);
		Common.appendToFile(datetime, test_id, type, "proc", filesize, t2, infile);
		Common.appendToFile(datetime, test_id, type, "save", filesize, t3, infile);
		Common.appendToFile(datetime, test_id, type, "sums", filesize, t1 + t2 + t3, infile);

		// 加入延时，保证后续的时间不一样。
		Thread.sleep(1);
	}

	// Cipher.ENCRYPT_MODE == 1
	// Cipher.DECRYPT_MODE == 2
	public Cipher getCipher(int type) throws Exception {
		Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
		SecretKeySpec keySpec = new SecretKeySpec(secKey.getEncoded(), "AES");
		IvParameterSpec ivSpec = new IvParameterSpec(IV);
		cipher.init(type, keySpec, ivSpec);
		return cipher;
	}

}
